#include "infer_demo.h"

void performance_test(ai::ModelManager* model, std::string file_path, std::vector<ai::InitModelData> init_model_datas, int max_channel_size)
{
  int ret = 0;
  stream::ImageBlob* blob = new stream::ImageBlob(stream::ImageBlobMode_BGR);
  cv::Mat pic = cv::imread(file_path);
  // std::cout << "pic:" << pic.cols << " " << pic.rows << std::endl;

  std::vector<int> fps_cnt(max_channel_size, 0);
  std::vector<float> npu_infer_fps(max_channel_size, 0);
  int64_t start_time = 0, end_time = 0;
  start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

  std::vector<std::string> model_ids;
  for (auto& imd : init_model_datas) {model_ids.push_back(imd.model_id);}
  // std::cout << "model_ids:" << model_ids.size() << std::endl;
  // std::cout << "model_ids:" << model_ids[0] << std::endl;

  std::vector<std::vector<int>> infer_idxs;
  // 模拟通道号
  int channel_id = 0;
  while(1)
  {
    
    // 喂数据
    std::vector<int> infer_ids;
    blob->img = pic;
    blob->source_id = channel_id;
    ret = model->set_data(model_ids, infer_ids, *blob);
    // std::cout << "set_data ret:" << ret << " infer_ids:" << infer_ids[0] << std::endl;
    // 数据设置成功 此数据必然会被拿去推理
    if (ret == 0) {
      channel_id++;
      if (channel_id >= max_channel_size) {channel_id = 0;}
      infer_idxs.push_back(infer_ids);
    }
    
    // 拿数据
    if (!infer_idxs.empty()) {
      infer_ids = infer_idxs.front();
      ai::InferResult infer_res;
      ret = model->get_data(model_ids, infer_ids, infer_res);
      // std::cout << "get_data ret:" << ret << " infer_ids:" << infer_ids[0] << std::endl;
      
      if (ret == 0) {
        infer_idxs.erase(infer_idxs.begin());
        fps_cnt[infer_res.source_id]++;
        for (auto& det_info : infer_res.infer_res) {delete det_info;}
      }
    }

    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (end_time - start_time > 5 * 1e3) {
      for (int idx = 0; idx< max_channel_size; idx++) {
        npu_infer_fps[idx] = 1.0 / ((end_time - start_time) / (fps_cnt[idx] * 1e3));
        std::cout << "channel:"<< idx << " infer_fps:" << npu_infer_fps[idx] << std::endl;
        fps_cnt[idx] = 0;
      }
      start_time = end_time;
    }

    // std::this_thread::sleep_for(milliseconds(3));

  }
  
}

int infer_det_img(ai::ModelManager* model, stream::CVOsdText* osd, std::vector<ai::InitModelData>& init_model_datas, std::string file_path) {
  int ret = 0;

  std::vector< bf::path > all_files;
  if (bf::is_directory(file_path)) {
    get_files_in_dir_recursive(file_path, all_files, {".jpg", ".png"});
  } else {
    all_files.push_back(bf::path(file_path));
  }

  std::vector<std::string> model_ids;
  for (auto& imd : init_model_datas) {model_ids.push_back(imd.model_id);}

  int frame_cnt = 0;
  std::vector<int> infer_ids;
  int plate_rec_idx = 1;
  stream::ImageBlob data_blob(stream::ImageBlobMode_BGR);
  for (auto& img_path : all_files) {
    

    cv::Mat car_img = cv::imread(img_path.string());
    // cv::Mat draw_img = car_img.clone();

    data_blob.source_id = frame_cnt;
    data_blob.img = car_img;
    ret = model->set_data(model_ids, infer_ids, data_blob);
    // std::cout << "infer_id:" << infer_id << std::endl;
    if (ret != 0) {continue;}

    ai::InferResult infer_res;
    while ((ret = model->get_data(model_ids, infer_ids, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }
    std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;

      cv::rectangle(data_blob.img, di->det_box, cv::Scalar(0, 255, 0), 2);
      cv::putText(data_blob.img, di->category+std::to_string(di->score).substr(0,4), cv::Point(di->det_box.x, di->det_box.y-10), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 3);
      // cv::imwrite(std::to_string(frame_cnt) + ".plate.jpg", plate_img);
      delete di;
    }
    frame_cnt++;
    if (infer_res.infer_res.size() > 0) {
      cv::imwrite(img_path.string() + ".draw.jpg", data_blob.img);
    }
  }
  return 0;
}


int infer_det_img_multi(ai::ModelManager* model, stream::CVOsdText* osd, std::vector<ai::InitModelData>& init_model_datas, std::string file_path) {
  int ret = 0;

  std::vector< bf::path > all_files;
  if (bf::is_directory(file_path)) {
    get_files_in_dir_recursive(file_path, all_files, {".jpg", ".png"});
  } else {
    all_files.push_back(bf::path(file_path));
  }

  int frame_cnt = 0;
  int plate_rec_idx = 1;
  stream::ImageBlob car_blob(stream::ImageBlobMode_BGR);
  for (auto& img_path : all_files) {
    cv::Mat car_img = cv::imread(img_path.string());

    car_blob.source_id = frame_cnt;
    car_blob.img = car_img;

    std::vector<std::string> model_ids;
    for (auto& imd : init_model_datas) {model_ids.push_back(imd.model_id);}
    std::vector<int> infer_ids(model_ids.size(), -1);
    
    ret = model->set_data(model_ids, infer_ids, car_blob);
    std::cout << "set_data ret:" << ret << std::endl;
    if (ret != 0) {continue;}

    std::string log_str = "infer_ids:\n";
    for (auto& ii : infer_ids) { log_str += std::to_string(ii) + ", "; }
    std::cout << log_str << std::endl;
    
    ai::InferResult infer_res;
    while ((ret = model->get_data(model_ids, infer_ids, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }
    std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;

      cv::rectangle(car_blob.img, di->det_box, cv::Scalar(0, 255, 0), 2);
      cv::putText(car_blob.img, di->category+std::to_string(di->score).substr(0,4), cv::Point(di->det_box.x, di->det_box.y-10), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 3);
      // cv::imwrite(std::to_string(frame_cnt) + ".plate.jpg", plate_img);
      delete di;
    }
    frame_cnt++;
    if (infer_res.infer_res.size() > 0) {
      cv::imwrite(img_path.string() + ".draw.jpg", car_blob.img);
    }
  }
  return 0;
}

int infer_det_mp4(ai::ModelManager* model, stream::CVOsdText* osd, std::vector<ai::InitModelData>& init_model_datas, std::string mp4_path) {
  int ret = 0;
  cv::VideoCapture reader = cv::VideoCapture(mp4_path);
  if (!reader.isOpened()) { 
    printf("ERROR. infer_mp4 !reader->isOpened() mp4_path: %s\n", mp4_path);
    return -1;
  }
  
  int width = reader.get(cv::CAP_PROP_FRAME_WIDTH);
  int height = reader.get(cv::CAP_PROP_FRAME_HEIGHT);
  double fps = reader.get(cv::CAP_PROP_FPS);
  if (fps < 1 || fps > 50) { fps = 15.;}
  printf("infer_mp4 width: %d, height: %d, fps: %f\n", width, height, fps);

  std::string writer_path = mp4_path + ".rec.mp4";
  cv::VideoWriter writer = cv::VideoWriter(writer_path, cv::VideoWriter::fourcc('a','v','c','1'), fps, cv::Size(width, height), true);

  int frame_cnt = 0;
  int infer_id = 0;
  stream::ImageBlob blob(stream::ImageBlobMode_BGR);
  blob.model_id = init_model_datas[0].model_id;
  while(true) {
    // if (frame_cnt > 500) {break;}
    
    bool ok = reader.read(blob.img);
    if (!ok) {break;}

    blob.source_id = frame_cnt++;
    model->set_data(blob.model_id, infer_id, blob);
    if (infer_id < 0) {continue;}

    ai::InferResult infer_res;
    while ((ret = model->get_data(blob.model_id, infer_id, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }

    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      cv::rectangle(blob.img, di->det_box, cv::Scalar(0, 0, 255), 2);
      cv::putText(blob.img, di->category + ":" + std::to_string(di->score).substr(0,4), cv::Point(di->det_box.x, di->det_box.y - 20), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 2);
      delete det_info;
    }

    writer.write(blob.img); 
    printf("frame_cnt:%d\n", frame_cnt);
  }
  reader.release();
  writer.release();
  
  return 0;
}

int infer_det_img_api(ai::ModelManager* model, std::vector<ai::InitModelData>& init_model_datas, std::string input_str, std::string& ret_str) {
  int ret = 0;

  Json::Reader reader;
  Json::Value root;
  if (!reader.parse(input_str, root)) { return json_read_parse_error; }

  std::string b64_img = root["b64Img"].asString();
  cv::Mat img;
  ret = cv_image_decode(b64_img, img, cv::IMREAD_COLOR);
  if (img.empty()) { 
    std::cout << "rgm_check: img is empty..." << std::endl; 
    return -1;
  }


  static int frame_cnt = 0;
  int infer_id = 0;
  int plate_rec_idx = 1;
  stream::ImageBlob car_blob(stream::ImageBlobMode_BGR);

    
  
  Json::Value ret_json;
  // cv::Mat draw_img = car_img.clone();

  car_blob.source_id = frame_cnt;
  car_blob.img = img;
  car_blob.model_id = init_model_datas[0].model_id;
  model->set_data(car_blob.model_id, infer_id, car_blob);
  std::cout << "infer_id:" << infer_id << std::endl;
  if (infer_id < 0) {return -1;}

  ai::InferResult infer_res;
  while ((ret = model->get_data(car_blob.model_id, infer_id, infer_res)) != 0) {
    std::this_thread::sleep_for(milliseconds(3));
  }
  std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

  ret_json["items"] = Json::arrayValue;
  for (auto& det_info : infer_res.infer_res) {
    ai::DetInfo* di = (ai::DetInfo*)det_info;
    // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;
    Json::Value item;
    item["score"] = di->score;
    item["x"] = di->det_box.x;
    item["y"] = di->det_box.y;
    item["w"] = di->det_box.width;
    item["h"] = di->det_box.height;
    item["category"] = di->category;
    ret_json["items"].append(item);

    delete di;
  }
  ret_json["code"] = 0;
  
  frame_cnt++;

  ret_str = ret_json.toStyledString();
  return 0;

}


int infer_det_img_to_via_file(ai::ModelManager* model, stream::CVOsdText* osd, std::vector<ai::InitModelData>& init_model_datas, std::string mp4_path) {
  int ret = 0;

  std::vector<boost::filesystem::path> last_dirs;
  get_last_dirs(mp4_path, last_dirs);

  int frame_cnt = 0;
  int infer_id = 0;
  int plate_rec_idx = 1;
  Json::Value root;
  stream::ImageBlob car_blob(stream::ImageBlobMode_BGR);
  for (auto& last_dir : last_dirs) {
    std::cout << "last_dir.string():" << last_dir.string() << std::endl;
    std::vector< bf::path > all_files;
    if (bf::is_directory(last_dir.string())) {
      get_files_in_dir_recursive(last_dir.string(), all_files, {".jpg", ".png"});
    } else {
      all_files.push_back(bf::path(last_dir));
    }
    
    for (auto& img_path : all_files) {
      
      Json::Value item;
      item["filename"] = bf::path(img_path).filename().string();
      cv::Mat car_img = cv::imread(img_path.string());
      // cv::Mat draw_img = car_img.clone();

      car_blob.source_id = frame_cnt;
      car_blob.img = car_img;
      car_blob.model_id = init_model_datas[0].model_id;
      model->set_data(car_blob.model_id, infer_id, car_blob);
      std::cout << "infer_id:" << infer_id << std::endl;
      if (infer_id < 0) {continue;}

      ai::InferResult infer_res;
      while ((ret = model->get_data(car_blob.model_id, infer_id, infer_res)) != 0) {
        std::this_thread::sleep_for(milliseconds(3));
      }
      std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

      Json::Value regions = Json::arrayValue;
      for (auto& det_info : infer_res.infer_res) {
        ai::DetInfo* di = (ai::DetInfo*)det_info;
        // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;
        if (di->category == "person") {
          Json::Value region;
          region["shape_attributes"]["name"] = "rect";
          region["shape_attributes"]["x"] = di->det_box.x;
          region["shape_attributes"]["y"] = di->det_box.y;
          region["shape_attributes"]["width"] = di->det_box.width;
          region["shape_attributes"]["height"] = di->det_box.height;
          region["region_attributes"]["label"] = di->category;
          regions.append(region);
        }

        delete di;
      }

      item["regions"] = regions;
      root[bf::path(img_path).filename().string()] = item;
      
      frame_cnt++;
    }

    std::ofstream imageyuv(last_dir.string()+"/via_region_data.json", std::ios::binary);
    if(imageyuv.is_open()) {
      std::string save_msg = root.toStyledString();
      imageyuv.write(save_msg.data(), save_msg.size());
      imageyuv.close();
    }
  }
  
  return 0;
}


int infer_cls_video(ai::ModelManager* model, stream::CVOsdText* osd, std::vector<ai::InitModelData>& init_model_datas, std::string file_path) {
  int ret = 0;

  std::vector< bf::path > all_files;
  if (bf::is_directory(file_path)) {
    get_files_in_dir_recursive(file_path, all_files, {".jpg", ".png"});
  } else {
    all_files.push_back(bf::path(file_path));
  }
  
  std::vector<std::string> all_img_paths;
  for (auto& img_path : all_files) { all_img_paths.push_back(img_path.string());}
  std::sort(all_img_paths.begin(), all_img_paths.end(), [](std::string& a, std::string& b){return a < b;});

  int frame_cnt = 0;
  int infer_id = 0;
  int plate_rec_idx = 1;
  stream::ImageBlob video_blob(stream::ImageBlobMode_BGR);
  // 填充8张图片
  for (int idx = 0; idx < 8 && idx < all_img_paths.size(); idx++) {
    cv::Mat img = cv::imread(all_img_paths[idx]);
    video_blob.imgs.push_back(img);
  }

  video_blob.source_id = frame_cnt;
  video_blob.model_id = init_model_datas[0].model_id;
  model->set_data(video_blob.model_id, infer_id, video_blob);
  std::cout << "infer_id:" << infer_id << std::endl;
  if (infer_id < 0) {return 0;}

  ai::InferResult infer_res;
  while ((ret = model->get_data(video_blob.model_id, infer_id, infer_res)) != 0) {
    std::this_thread::sleep_for(milliseconds(3));
  }
  std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

  for (auto& cls_info : infer_res.infer_res) {
    ai::ClsInfo* cls = (ai::ClsInfo*)cls_info;
    for (int cls_idx = 0; cls_idx < cls->top5_score.size(); cls_idx++) {
      std::cout << "class:" << cls->top5_class[cls_idx] << ", score:" << cls->top5_score[cls_idx] << std::endl;
    }
    delete cls;
  }
  frame_cnt++;

  return 0;
}



int infer_face_img_det_rec(ai::ModelManager* model, stream::CVOsdText* osd, std::string file_path) {
  int ret = 0;

  std::vector< bf::path > all_files;
  if (bf::is_directory(file_path)) {
    get_files_in_dir_recursive(file_path, all_files, {".jpg", ".png"});
  } else {
    all_files.push_back(bf::path(file_path));
  }
  // std::vector<std::pair<int64_t, std::string>> img_paths;
  // for (auto& f : all_files) { 
  //   std::string filename = f.filename().string();
  //   std::vector<std::string> ips;
  //   boost::algorithm::split(ips, filename, boost::is_any_of("."), boost::token_compress_on);
  //   int64_t time_point = std::stoll(ips[1]);
  //   img_paths.push_back(std::make_pair(time_point, f.string()));
  // }
  // sort(img_paths.begin(), img_paths.end(), [](std::pair<int64_t, std::string> a, std::pair<int64_t, std::string> b){return a.first < b.first;});

  int frame_cnt = 0;
  int infer_id = 0;
  int face_infer_id = 1;
  int det_id = 0;
  int rec_id = 0;
  stream::ImageBlob det_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob rec_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob draw_data(stream::ImageBlobMode_BGR);
    // if (frame_cnt > 500) {break;}
  for (auto& img_path : all_files) {
    

    cv::Mat ori_img = cv::imread(img_path.string());
    cv::Mat draw_img = ori_img.clone();

    det_blob.img = ori_img;
    det_blob.source_id = det_id++;
    det_blob.model_id = "face";
    model->set_data(det_blob.model_id, infer_id, det_blob);
    std::cout << "infer_id:" << infer_id << std::endl;
    if (infer_id < 0) {return -1;}

    ai::InferResult infer_res;
    while ((ret = model->get_data(det_blob.model_id, infer_id, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }
    std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

    int plate_idx=0;
    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;

      // cv::rectangle(blob.img, di->det_box, cv::Scalar(0, 0, 255), 2);
      // cv::Mat face_img = ori_img(di->det_box).clone();
      cv::Mat align_face_img;
      ret = common::get_face_align_crop_img(ori_img, *di, align_face_img);
      
      frame_cnt++;
      // cv::imwrite(std::to_string(frame_cnt) + ".plate.jpg", plate_img);
      // cv::imwrite(std::to_string(frame_cnt) + ".car.jpg", car_img);


      rec_blob.img = align_face_img;
      rec_blob.source_id = rec_id++;
      rec_blob.model_id = "facerec";
      model->set_data(rec_blob.model_id, face_infer_id, rec_blob);

      ai::InferResult rec_infer_res;
      while ((ret = model->get_data(rec_blob.model_id, face_infer_id, rec_infer_res)) != 0) {
        std::this_thread::sleep_for(milliseconds(3));
      }

      if (rec_infer_res.infer_res.size() > 0) {
        ai::DetInfo* di2 = (ai::DetInfo*)rec_infer_res.infer_res[0];
        std::vector<float> embeddings=di2->kps;
        printf("face embeddings:\n");
        for(int i=0;i<embeddings.size();i++)
        {
          printf("%.4f\n",embeddings.data()[i]);
        }
        printf("\n");
        delete di2;
      }

      delete di;

    }

    // writer.write(draw_img); 
    // printf("frame_cnt:%d\n", frame_cnt);
  }


  
  return 0;
}
